!
!(c) Matthew Kennel, Institute for Nonlinear Science (2004)
!
! The KDTREE2 software is licensed under the terms of the Academic Free
! Software License, listed herein.  In addition, users of this software
! must give appropriate citation in relevant technical documentation or
! journal paper to the author, Matthew B. Kennel, Institute For
! Nonlinear Science, preferably via a reference to the www.arxiv.org
! repository of this document, {\tt www.arxiv.org e-print:
! physics/0408067}.  This requirement will be deemed to be advisory and
! not mandatory as is necessary to permit the free inclusion of the
! present software with any software licensed under the terms of any
! version of the GNU General Public License, or GNU Library General
! Public License.
!
! Academic Free License
! Version 1.1
!
! This Academic Free License applies to any original work of authorship
! (the "Original Work") whose owner (the "Licensor") has placed the
! following notice immediately following the copyright notice for the
! Original Work: "Licensed under the Academic Free License version 1.1."
!
! Grant of License. Licensor hereby grants to any person obtaining a
! copy of the Original Work ("You") a world-wide, royalty-free,
! non-exclusive, perpetual, non-sublicenseable license (1) to use, copy,
! modify, merge, publish, perform, distribute and/or sell copies of the
! Original Work and derivative works thereof, and (2) under patent
! claims owned or controlled by the Licensor that are embodied in the
! Original Work as furnished by the Licensor, to make, use, sell and
! offer for sale the Original Work and derivative works thereof, subject
! to the following conditions.
!
! Right of Attribution. Redistributions of the Original Work must
! reproduce all copyright notices in the Original Work as furnished by
! the Licensor, both in the Original Work itself and in any
! documentation and/or other materials provided with the distribution of
! the Original Work in executable form.
!
! Exclusions from License Grant. Neither the names of Licensor, nor the
! names of any contributors to the Original Work, nor any of their
! trademarks or service marks, may be used to endorse or promote
! products derived from this Original Work without express prior written
! permission of the Licensor.
!
! WARRANTY AND DISCLAIMERS. LICENSOR WARRANTS THAT THE COPYRIGHT IN AND
! TO THE ORIGINAL WORK IS OWNED BY THE LICENSOR OR THAT THE ORIGINAL
! WORK IS DISTRIBUTED BY LICENSOR UNDER A VALID CURRENT LICENSE FROM THE
! COPYRIGHT OWNER. EXCEPT AS EXPRESSLY STATED IN THE IMMEDIATELY
! PRECEEDING SENTENCE, THE ORIGINAL WORK IS PROVIDED UNDER THIS LICENSE
! ON AN "AS IS" BASIS, WITHOUT WARRANTY, EITHER EXPRESS OR IMPLIED,
! INCLUDING, WITHOUT LIMITATION, THE WARRANTY OF NON-INFRINGEMENT AND
! WARRANTIES THAT THE ORIGINAL WORK IS MERCHANTABLE OR FIT FOR A
! PARTICULAR PURPOSE. THE ENTIRE RISK AS TO THE QUALITY OF THE ORIGINAL
! WORK IS WITH YOU. THIS DISCLAIMER OF WARRANTY CONSTITUTES AN ESSENTIAL
! PART OF THIS LICENSE. NO LICENSE TO ORIGINAL WORK IS GRANTED HEREUNDER
! EXCEPT UNDER THIS DISCLAIMER.
!
! LIMITATION OF LIABILITY. UNDER NO CIRCUMSTANCES AND UNDER NO LEGAL
! THEORY, WHETHER TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE,
! SHALL THE LICENSOR BE LIABLE TO ANY PERSON FOR ANY DIRECT, INDIRECT,
! SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES OF ANY CHARACTER ARISING
! AS A RESULT OF THIS LICENSE OR THE USE OF THE ORIGINAL WORK INCLUDING,
! WITHOUT LIMITATION, DAMAGES FOR LOSS OF GOODWILL, WORK STOPPAGE,
! COMPUTER FAILURE OR MALFUNCTION, OR ANY AND ALL OTHER COMMERCIAL
! DAMAGES OR LOSSES, EVEN IF SUCH PERSON SHALL HAVE BEEN INFORMED OF THE
! POSSIBILITY OF SUCH DAMAGES. THIS LIMITATION OF LIABILITY SHALL NOT
! APPLY TO LIABILITY FOR DEATH OR PERSONAL INJURY RESULTING FROM SUCH
! PARTY'S NEGLIGENCE TO THE EXTENT APPLICABLE LAW PROHIBITS SUCH
! LIMITATION. SOME JURISDICTIONS DO NOT ALLOW THE EXCLUSION OR
! LIMITATION OF INCIDENTAL OR CONSEQUENTIAL DAMAGES, SO THIS EXCLUSION
! AND LIMITATION MAY NOT APPLY TO YOU.
!
! License to Source Code. The term "Source Code" means the preferred
! form of the Original Work for making modifications to it and all
! available documentation describing how to access and modify the
! Original Work. Licensor hereby agrees to provide a machine-readable
! copy of the Source Code of the Original Work along with each copy of
! the Original Work that Licensor distributes. Licensor reserves the
! right to satisfy this obligation by placing a machine-readable copy of
! the Source Code in an information repository reasonably calculated to
! permit inexpensive and convenient access by You for as long as
! Licensor continues to distribute the Original Work, and by publishing
! the address of that information repository in a notice immediately
! following the copyright notice that applies to the Original Work.
!
! Mutual Termination for Patent Action. This License shall terminate
! automatically and You may no longer exercise any of the rights granted
! to You by this License if You file a lawsuit in any court alleging
! that any OSI Certified open source software that is licensed under any
! license containing this "Mutual Termination for Patent Action" clause
! infringes any patent claims that are essential to use that software.
!
! This license is Copyright (C) 2002 Lawrence E. Rosen. All rights
! reserved. Permission is hereby granted to copy and distribute this
! license without modification. This license may not be modified without
! the express written permission of its copyright owner.
!

module mod_kdtree2

  use mod_realkinds
  use mod_stdio
  use mod_message
  use mod_kdtree2_priority

  implicit none

  private

  integer, parameter :: sp = rk4
  integer, parameter :: dp = rk8
  integer, parameter :: kdkind = rk8
  !
  ! K-D tree routines in Fortran 90 by Matt Kennel.
  ! Original program was written in Sather by Steve Omohundro and
  ! Matt Kennel.  Only the Euclidean metric is supported.
  !
  ! This module is identical to 'kd_tree', except that the order
  ! of subscripts is reversed in the data file.
  ! In otherwords for an embedding of N D-dimensional vectors, the
  ! data file is here, in natural Fortran order  data(1:D, 1:N)
  ! because Fortran lays out columns first,
  !
  ! whereas conventionally (C-style) it is data(1:N,1:D)
  ! as in the original kd_tree module.
  !
  !-------------DATA TYPE, CREATION, DELETION---------------------
  public :: kdkind
  public :: kdtree2, kdtree2_result
  public :: tree_node, kdtree2_create, kdtree2_destroy
  !---------------------------------------------------------------
  !-------------------SEARCH ROUTINES-----------------------------
  public :: kdtree2_n_nearest, kdtree2_n_nearest_around_point
  ! Return fixed number of nearest neighbors around arbitrary vector,
  ! or extant point in dataset, with decorrelation window.
  !
  public :: kdtree2_r_nearest, kdtree2_r_nearest_around_point
  ! Return points within a fixed ball of arb vector/extant point
  !
  public :: kdtree2_sort_results
  ! Sort, in order of increasing distance, rseults from above.
  !
  public :: kdtree2_r_count, kdtree2_r_count_around_point
  ! Count points within a fixed ball of arb vector/extant point
  !
  public :: kdtree2_n_nearest_brute_force, kdtree2_r_nearest_brute_force
  ! brute force of kdtree2_[n|r]_nearest
  !----------------------------------------------------------------

  integer, parameter :: bucket_size = 13

  ! The maximum number of points to keep in a terminal node.

  type interval
    real(kdkind) :: lower
    real(kdkind) :: upper
  end type interval

  type :: tree_node
    ! an internal tree node
    private
    integer :: cut_dim
    ! the dimension to cut
    real(kdkind) :: cut_val
    ! where to cut the dimension
    real(kdkind) :: cut_val_left
    real(kdkind) :: cut_val_right
    ! improved cutoffs knowing the spread in child boxes.
    integer :: l
    integer :: u
    type(tree_node), pointer :: left
    type(tree_node), pointer :: right
    type(interval), pointer :: box(:) => null()
    ! child pointers
    ! Points included in this node are indexes[k] with k \in [l,u]
  end type tree_node

  type :: kdtree2
    ! Global information about the tree, one per tree
    integer :: dimen = 0
    integer :: n = 0
    ! dimensionality and total # of points
    real(kdkind), pointer, contiguous :: the_data(:,:) => null()
    ! pointer to the actual data array
    !
    !  IMPORTANT NOTE:  IT IS DIMENSIONED   the_data(1:d,1:N)
    !  which may be opposite of what may be conventional.
    !  This is, because in Fortran, the memory layout is such that
    !  the first dimension is in sequential order.  Hence, with
    !  (1:d,1:N), all components of the vector will be in consecutive
    !  memory locations.  The search time is dominated by the
    !  evaluation of distances in the terminal nodes.  Putting all
    !  vector components in consecutive memory location improves
    !  memory cache locality, and hence search speed, and may enable
    !  vectorization on some processors and compilers.

    integer, pointer :: ind(:) => null()
    ! permuted index into the data, so that indexes[l..u] of some
    ! bucket represent the indexes of the actual points in that
    ! bucket.
    logical :: sort = .false.
    ! do we always sort output results?
    logical :: rearrange = .false.
    real(kdkind), pointer, contiguous :: rearranged_data(:,:) => null()
    ! if (rearrange .eqv. .true.) then rearranged_data has been
    ! created so that rearranged_data(:,i) = the_data(:,ind(i)),
    ! permitting search to use more cache-friendly rearranged_data, at
    ! some initial computation and storage cost.
    type(tree_node), pointer :: root => null()
    ! root pointer of the tree
  end type kdtree2

  type :: tree_search_record
    !
    ! One of these is created for each search.
    !
    private
    !
    ! Many fields are copied from the tree structure, in order to
    ! speed up the search.
    !
    integer :: dimen
    integer :: nn
    integer :: nfound
    real(kdkind) :: ballsize
    integer :: centeridx = 999
    integer :: correltime = 9999
    ! exclude points within 'correltime' of 'centeridx', iff centeridx >= 0
    integer :: nalloc     ! how much allocated for results(:)?
    logical :: rearrange  ! are the data rearranged or original?
    ! did the # of points found overflow the storage provided?
    logical :: overflow
    real(kdkind), pointer :: qv(:)  ! query vector
    type(kdtree2_result), pointer :: results(:) ! results
    type(pq) :: pq
    real(kdkind), pointer, contiguous :: rdata(:,:)  ! temp pointer to data
    integer, pointer :: ind(:)     ! temp pointer to indexes
  end type tree_search_record

  ! everything else is private.

  ! A GLOBAL VARIABLE for search
  type(tree_search_record), save, target :: sr

  contains

  function kdtree2_create(input_data,indim,balanced,sort,rearrange) result (mr)
    implicit none
    !
    ! create the actual tree structure, given an input array of data.
    !
    ! Note, input data is input_data(1:d,1:N), NOT the other way around.
    ! THIS IS THE REVERSE OF THE PREVIOUS VERSION OF THIS MODULE.
    ! The reason for it is cache friendliness, improving performance.
    !
    ! Optional arguments:  If 'dim' is specified, then the tree
    !                      will only search the first 'dim' components
    !                      of input_data, otherwise, dim is inferred
    !                      from SIZE(input_data,1).
    !
    !                      if balanced .eqv. .true. then the tree
    !                      buildin select exact median to have fully
    !                      balanced tree.
    !
    !                      if sort .eqv. .true. then output results
    !                      will be sorted by increasing distance.
    !                      default=.false., as it is faster to not sort.
    !
    !                      if rearrange .eqv. .true. then an internal
    !                      copy of the data, rearranged by terminal node,
    !                      will be made for cache friendliness.
    !                      default=.true., as it speeds searches, but
    !                      building takes longer, and extra memory is used.
    !
    ! .. Function Return Cut_value ..
    type(kdtree2), pointer :: mr
    integer, intent(in), optional :: indim
    logical, intent(in), optional :: balanced
    logical, intent(in), optional :: sort
    logical, intent(in), optional :: rearrange
    ! ..
    ! .. Array Arguments ..
    real(kdkind), target, dimension(:,:) :: input_data
    !
    integer :: i
    !
    logical :: lbal
    ! ..
    allocate (mr)
    mr%the_data => input_data
    ! pointer assignment

    if ( present(indim) ) then
      mr%dimen = indim
    else
      mr%dimen = size(input_data,1)
    end if
    mr%n = size(input_data,2)
    lbal = .false.
    if ( present(balanced) ) then
      lbal = .true.
    end if

    if ( mr%dimen > mr%n ) then
      !  unlikely to be correct
      write(stderr,*) 'KD_TREE_TRANS: likely user error.'
      write(stderr,*) 'KD_TREE_TRANS: You passed in matrix with D=',mr%dimen
      write(stderr,*) 'KD_TREE_TRANS: and N=',mr%n
      write(stderr,*) 'KD_TREE_TRANS: note, that new format is data(1:D,1:N)'
      write(stderr,*) 'KD_TREE_TRANS: with usually N >> D.'
      write(stderr,*) &
        'If N =approx= D, then a k-d tree is not an appropriate data structure.'
      call die('KD_TREE_TRANS','not appropriate',1)
    end if

    call build_tree(mr,lbal)

    if ( present(sort) ) then
      mr%sort = sort
    else
      mr%sort = .false.
    end if

    if ( present(rearrange) ) then
      mr%rearrange = rearrange
    else
      mr%rearrange = .false.
    end if

    if ( mr%rearrange ) then
      allocate(mr%rearranged_data(mr%dimen,mr%n))
      do i = 1, mr%n
        mr%rearranged_data(:,i) = mr%the_data(:,mr%ind(i))
      end do
    else
      nullify(mr%rearranged_data)
    endif
  end function kdtree2_create

  subroutine build_tree(tp,lbal)
    implicit none
    type (kdtree2), pointer :: tp
    logical :: lbal
    ! ..
    integer :: j
    type(tree_node), pointer :: dummy => null()
    ! ..
    allocate(tp%ind(tp%n))
    do concurrent ( j = 1:tp%n )
      tp%ind(j) = j
    end do
    tp%root => build_tree_for_range(tp,1,tp%n,dummy,lbal)
  end subroutine build_tree

  recursive function build_tree_for_range(tp,l,u,parent,lbal) result (res)
    implicit none
    ! .. Function Return Cut_value ..
    type (tree_node), pointer :: res
    ! ..
    ! .. Structure Arguments ..
    type (kdtree2), pointer :: tp
    type (tree_node), pointer :: parent
    ! ..
    ! .. Scalar Arguments ..
    integer, intent (in) :: l, u
    logical, intent (in) :: lbal
    ! ..
    ! .. Local Scalars ..
    integer :: i, c, m, dimen
    logical :: recompute
    real(kdkind) :: average

    ! first compute min and max
    dimen = tp%dimen

    if ( u < l ) then
      ! no points in this box
      nullify(res)
      return
    end if

    allocate(res)
    allocate(res%box(dimen))

    ! First, compute an APPROXIMATE bounding box of all points
    ! associated with this node.

    if ( (u-l) <= bucket_size ) then
      !
      ! always compute true bounding box for terminal nodes.
      !
      do i = 1, dimen
        call spread_in_coordinate(tp,i,l,u,res%box(i))
      end do
      res%cut_dim = 0
      res%cut_val = 0.0
      res%l = l
      res%u = u
      res%left => null()
      res%right => null()
    else
      !
      ! modify approximate bounding box.  This will be an
      ! overestimate of the true bounding box, as we are only recomputing
      ! the bounding box for the dimension that the parent split on.
      !
      ! Going to a true bounding box computation would significantly
      ! increase the time necessary to build the tree, and usually
      ! has only a very small difference.  This box is not used
      ! for searching but only for deciding which coordinate to split on.
      !
      do i = 1, dimen
        recompute = .true.
        if ( associated(parent) ) then
          if ( i /= parent%cut_dim ) then
            recompute = .false.
          end if
        end if
        if ( recompute ) then
          call spread_in_coordinate(tp,i,l,u,res%box(i))
        else
          res%box(i) = parent%box(i)
        end if
      end do

      c = maxloc(res%box(1:dimen)%upper-res%box(1:dimen)%lower,1)

      !
      ! c is the identity of which coordinate has the greatest spread.
      !

      if ( lbal ) then
        ! select exact median to have fully balanced tree.
        m = (l+u)/2
        call select_on_coordinate(tp%the_data,tp%ind,c,m,l,u)
      else
        !
        ! select point halfway between min and max, as per A. Moore,
        ! who says this helps in some degenerate cases, or
        ! actual arithmetic average.
        !
        if ( .true. ) then
          ! actually compute average
          average = sum(tp%the_data(c,tp%ind(l:u))) / real(u-l+1,kdkind)
        else
          average = (res%box(c)%upper + res%box(c)%lower)/2.0
        end if

        res%cut_val = average
        m = select_on_coordinate_value(tp%the_data,tp%ind,c,average,l,u)
      end if

      ! moves indexes around
      res%cut_dim = c
      res%l = l
      res%u = u

      res%left => build_tree_for_range(tp,l,m,res,lbal)
      res%right => build_tree_for_range(tp,m+1,u,res,lbal)

      if ( associated(res%right) .eqv. .false. ) then
        res%box = res%left%box
        res%cut_val_left = res%left%box(c)%upper
        res%cut_val = res%cut_val_left
      else if ( associated(res%left) .eqv. .false. ) then
        res%box = res%right%box
        res%cut_val_right = res%right%box(c)%lower
        res%cut_val = res%cut_val_right
      else
        res%cut_val_right = res%right%box(c)%lower
        res%cut_val_left = res%left%box(c)%upper
        res%cut_val = (res%cut_val_left + res%cut_val_right)/2

        ! now remake the true bounding box for self.
        ! Since we are taking unions (in effect) of a tree structure,
        ! this is much faster than doing an exhaustive
        ! search over all points
        res%box%upper = max(res%left%box%upper,res%right%box%upper)
        res%box%lower = min(res%left%box%lower,res%right%box%lower)
      end if
    end if

    contains

    integer function select_on_coordinate_value(v,ind,c,alpha,li,ui) result(res)
      implicit none
      ! Move elts of ind around between l and u, so that all points
      ! <= than alpha (in c cooordinate) are first, and then
      ! all points > alpha are second.

      !
      ! Algorithm (matt kennel).
      !
      ! Consider the list as having three parts: on the left,
      ! the points known to be <= alpha.  On the right, the points
      ! known to be > alpha, and in the middle, the currently unknown
      ! points.   The algorithm is to scan the unknown points, starting
      ! from the left, and swapping them so that they are added to
      ! the left stack or the right stack, as appropriate.
      !
      ! The algorithm finishes when the unknown stack is empty.
      !
      ! .. Scalar Arguments ..
      integer, intent (in) :: c, li, ui
      real(kdkind), intent(in) :: alpha
      ! ..
      real(kdkind), dimension(1:,1:) :: v
      integer, dimension(1:) :: ind
      integer :: tmp
      ! ..
      integer :: lb, rb
      !
      ! The points known to be <= alpha are in
      ! [l,lb-1]
      !
      ! The points known to be > alpha are in
      ! [rb+1,u].
      !
      ! Therefore we add new points into lb or
      ! rb as appropriate.  When lb=rb
      ! we are done.  We return the location of the last point <= alpha.
      !
      !
      lb = li
      rb = ui

      do while ( lb < rb )
        if ( v(c,ind(lb)) <= alpha ) then
          ! it is good where it is.
          lb = lb+1
        else
          ! swap it with rb.
          tmp = ind(lb)
          ind(lb) = ind(rb)
          ind(rb) = tmp
          rb = rb-1
        end if
      end do

      ! now lb .eq. ub
      if ( v(c,ind(lb)) <= alpha ) then
        res = lb
      else
        res = lb-1
      end if

    end function select_on_coordinate_value

    subroutine select_on_coordinate(v,ind,c,k,li,ui)
      implicit none
      ! Move elts of ind around between l and u, so that the kth
      ! element
      ! is >= those below, <= those above, in the coordinate c.
      ! .. Scalar Arguments ..
      integer, intent (in) :: c, k, li, ui
      ! ..
      integer :: i, l, m, s, t, u
      ! ..
      real(kdkind), dimension(:,:) :: v
      integer, dimension(:) :: ind
      ! ..
      l = li
      u = ui
      do while ( l < u )
        t = ind(l)
        m = l
        do i = l + 1, u
          if ( v(c,ind(i)) < v(c,t) ) then
            m = m + 1
            s = ind(m)
            ind(m) = ind(i)
            ind(i) = s
          end if
        end do
        s = ind(l)
        ind(l) = ind(m)
        ind(m) = s
        if ( m <= k ) l = m + 1
        if ( m >= k ) u = m - 1
      end do
    end subroutine select_on_coordinate

  end function build_tree_for_range

  subroutine spread_in_coordinate(tp,c,l,u,interv)
    implicit none
    ! the spread in coordinate 'c', between l and u.
    !
    ! Return lower bound in 'smin', and upper in 'smax',
    ! ..
    ! .. Structure Arguments ..
    type (kdtree2), pointer :: tp
    type(interval), intent(out) :: interv
    ! ..
    ! .. Scalar Arguments ..
    integer, intent(in) :: c, l, u
    ! ..
    ! .. Local Scalars ..
    real(kdkind) :: last, lmax, lmin, t, smin, smax
    integer :: i, ulocal
    ! ..
    ! .. Local Arrays ..
    real(kdkind), pointer, contiguous, dimension(:,:) :: v
    integer, pointer, dimension(:) :: ind
    ! ..
    v => tp%the_data
    ind => tp%ind
    smin = v(c,ind(l))
    smax = smin

    ulocal = u

    do i = l + 2, ulocal, 2
      lmin = v(c,ind(i-1))
      lmax = v(c,ind(i))
      if ( lmin > lmax ) then
        t = lmin
        lmin = lmax
        lmax = t
      end if
      if ( smin > lmin ) smin = lmin
      if ( smax < lmax ) smax = lmax
    end do
    if ( i == ulocal+1 ) then
      last = v(c,ind(ulocal))
      if ( smin > last ) smin = last
      if ( smax < last ) smax = last
    end if

    interv%lower = smin
    interv%upper = smax
  end subroutine spread_in_coordinate

  subroutine kdtree2_destroy(tp)
    implicit none
    ! Deallocates all memory for the tree, except input data matrix
    ! .. Structure Arguments ..
    type (kdtree2), pointer :: tp
    ! ..
    call destroy_node(tp%root)

    deallocate(tp%ind)
    nullify(tp%ind)

    if ( tp%rearrange ) then
      deallocate(tp%rearranged_data)
      nullify(tp%rearranged_data)
    end if

    deallocate(tp)

    contains

    recursive subroutine destroy_node(np)
      implicit none
      ! .. Structure Arguments ..
      type (tree_node), pointer :: np
      ! ..
      ! ..
      if ( associated(np%left) ) then
        call destroy_node(np%left)
        nullify(np%left)
      end if
      if ( associated(np%right) ) then
        call destroy_node(np%right)
        nullify(np%right)
      end if
      if ( associated(np%box) ) deallocate(np%box)
      deallocate(np)
    end subroutine destroy_node

  end subroutine kdtree2_destroy

  subroutine kdtree2_n_nearest(tp,qv,nn,results)
    implicit none
    ! Find the 'nn' vectors in the tree nearest to 'qv' in euclidean norm
    ! returning their indexes and distances in 'indexes' and 'distances'
    ! arrays already allocated passed to this subroutine.
    type (kdtree2), pointer :: tp
    real(kdkind), target, dimension(:), intent (in) :: qv
    integer, intent(inout) :: nn
    type(kdtree2_result), dimension(:), target :: results

    sr%ballsize = huge(1.0)
    sr%qv => qv
    sr%nn = nn
    sr%nfound = 0
    sr%centeridx = -1
    sr%correltime = 0
    sr%overflow = .false.

    sr%results => results

    sr%nalloc = nn   ! will be checked

    sr%ind => tp%ind
    sr%rearrange = tp%rearrange
    if (tp%rearrange) then
      sr%rdata => tp%rearranged_data
    else
      sr%rdata => tp%the_data
    end if
    sr%dimen = tp%dimen

    call validate_query_storage(nn)
    sr%pq = pq_create(results)

    call search(tp%root)

    if (tp%sort) then
      call kdtree2_sort_results(nn, results)
    endif
    nn = sr%nfound
  end subroutine kdtree2_n_nearest

  subroutine kdtree2_n_nearest_around_point(tp,idxin,correltime,nn,results)
    implicit none
    ! Find the 'nn' vectors in the tree nearest to point 'idxin',
    ! with correlation window 'correltime', returing results in
    ! results(:), which must be pre-allocated upon entry.
    type (kdtree2), pointer :: tp
    integer, intent (in) :: idxin, correltime
    integer, intent (inout) :: nn
    type(kdtree2_result), dimension(:), target :: results

    allocate (sr%qv(tp%dimen))
    sr%qv = tp%the_data(:,idxin) ! copy the vector
    sr%ballsize = huge(1.0)      ! the largest real(kdkind) number
    sr%centeridx = idxin
    sr%correltime = correltime

    sr%nn = nn
    sr%nfound = 0

    sr%dimen = tp%dimen
    sr%nalloc = nn

    sr%results => results

    sr%ind => tp%ind
    sr%rearrange = tp%rearrange

    if (sr%rearrange) then
      sr%rdata => tp%rearranged_data
    else
      sr%rdata => tp%the_data
    end if

    call validate_query_storage(nn)
    sr%pq = pq_create(results)

    call search(tp%root)

    if (tp%sort) then
      call kdtree2_sort_results(nn, results)
    endif
    deallocate (sr%qv)
    nn = sr%nfound
  end subroutine kdtree2_n_nearest_around_point

  subroutine kdtree2_r_nearest(tp,qv,r2,nfound,nalloc,results)
    implicit none
    ! find the nearest neighbors to point 'idxin', within SQUARED
    ! Euclidean distance 'r2'.   Upon ENTRY, nalloc must be the
    ! size of memory allocated for results(1:nalloc).  Upon
    ! EXIT, nfound is the number actually found within the ball.
    !
    !  Note that if nfound .gt. nalloc then more neighbors were found
    !  than there were storage to store.  The resulting list is NOT
    !  the smallest ball inside norm r^2
    !
    ! Results are NOT sorted unless tree was created with sort option.
    type(kdtree2), pointer :: tp
    real(kdkind), target, dimension(:), intent (in) :: qv
    real(kdkind), intent(in) :: r2
    integer, intent(out) :: nfound
    integer, intent(in) :: nalloc
    type(kdtree2_result), dimension(:), target :: results

    !
    sr%qv => qv
    sr%ballsize = r2
    sr%nn = 0      ! flag for fixed ball search
    sr%nfound = 0
    sr%centeridx = -1
    sr%correltime = 0

    sr%results => results

    call validate_query_storage(nalloc)
    sr%nalloc = nalloc
    sr%overflow = .false.
    sr%ind => tp%ind
    sr%rearrange= tp%rearrange

    if (tp%rearrange) then
      sr%rdata => tp%rearranged_data
    else
      sr%rdata => tp%the_data
    endif
    sr%dimen = tp%dimen

    !
    !sr%dsl = huge(sr%dsl)    ! set to huge positive values
    !sr%il = -1               ! set to invalid indexes
    !

    call search(tp%root)
    nfound = sr%nfound
    if ( tp%sort ) then
      call kdtree2_sort_results(nfound, results)
    endif

    if ( sr%overflow ) then
      write(stderr,*) 'KD_TREE_TRANS'
      write(stderr,*) &
        'warning! return from kdtree2_r_nearest found more neighbors'
      write(stderr,*) &
        'than storage was provided for.  Answer is NOT smallest ball'
      write(stderr,*) &
        'with that number of neighbors!  I.e. it is wrong.'
    end if
  end subroutine kdtree2_r_nearest

  subroutine kdtree2_r_nearest_around_point(tp,idxin, &
                          correltime,r2,nfound,nalloc,results)
    implicit none
    !
    ! Like kdtree2_r_nearest, but around a point 'idxin' already existing
    ! in the data set.
    !
    ! Results are NOT sorted unless tree was created with sort option.
    !
    type(kdtree2), pointer :: tp
    integer, intent(in) :: idxin, correltime, nalloc
    real(kdkind), intent(in) :: r2
    integer, intent(out) :: nfound
    type(kdtree2_result), dimension(:), target :: results
    !
    allocate (sr%qv(tp%dimen))
    sr%qv = tp%the_data(:,idxin) ! copy the vector
    sr%ballsize = r2
    sr%nn = 0    ! flag for fixed r search
    sr%nfound = 0
    sr%centeridx = idxin
    sr%correltime = correltime

    sr%results => results

    sr%nalloc = nalloc
    sr%overflow = .false.

    call validate_query_storage(nalloc)

    !    sr%dsl = HUGE(sr%dsl)    ! set to huge positive values
    !    sr%il = -1               ! set to invalid indexes

    sr%ind => tp%ind
    sr%rearrange = tp%rearrange

    if (tp%rearrange) then
      sr%rdata => tp%rearranged_data
    else
      sr%rdata => tp%the_data
    endif
    sr%rearrange = tp%rearrange
    sr%dimen = tp%dimen
    !
    !sr%dsl = huge(sr%dsl)    ! set to huge positive values
    !sr%il = -1               ! set to invalid indexes
    !
    call search(tp%root)
    nfound = sr%nfound
    if ( tp%sort ) then
      call kdtree2_sort_results(nfound,results)
    endif

    if ( sr%overflow ) then
      write (stderr,*) 'KD_TREE_TRANS'
      write (stderr,*) &
         'warning! return from kdtree2_r_nearest found more neighbors'
      write (stderr,*) &
        'than storage was provided for.  Answer is NOT smallest ball'
      write (stderr,*) &
        'with that number of neighbors!  I.e. it is wrong.'
    end if

    deallocate (sr%qv)
  end subroutine kdtree2_r_nearest_around_point

  integer function kdtree2_r_count(tp,qv,r2) result(nfound)
    implicit none
    ! Count the number of neighbors within square distance 'r2'.
    type (kdtree2), pointer :: tp
    real(kdkind), dimension(:), target, intent(in) :: qv
    real(kdkind), intent(in) :: r2
    !
    sr%qv => qv
    sr%ballsize = r2

    sr%nn = 0       ! flag for fixed r search
    sr%nfound = 0
    sr%centeridx = -1
    sr%correltime = 0

    nullify(sr%results) ! for some reason, FTN 95 chokes on '=> null()'

    sr%nalloc = 0            ! we do not allocate any storage but that's OK
                             ! for counting.
    sr%ind => tp%ind
    sr%rearrange = tp%rearrange
    if (tp%rearrange) then
      sr%rdata => tp%rearranged_data
    else
      sr%rdata => tp%the_data
    endif
    sr%dimen = tp%dimen
    !
    !sr%dsl = huge(sr%dsl)    ! set to huge positive values
    !sr%il = -1               ! set to invalid indexes
    !
    sr%overflow = .false.

    call search(tp%root)

    nfound = sr%nfound
  end function kdtree2_r_count

  integer function kdtree2_r_count_around_point(tp,idxin, &
                          correltime,r2) result(nfound)
    implicit none
    ! Count the number of neighbors within square distance 'r2' around
    ! point 'idxin' with decorrelation time 'correltime'.
    !
    type(kdtree2), pointer :: tp
    integer, intent(in) :: correltime, idxin
    real(kdkind), intent(in) :: r2
    !
    allocate (sr%qv(tp%dimen))
    sr%qv = tp%the_data(:,idxin)
    sr%ballsize = r2

    sr%nn = 0       ! flag for fixed r search
    sr%nfound = 0
    sr%centeridx = idxin
    sr%correltime = correltime
    nullify(sr%results)

    sr%nalloc = 0            ! we do not allocate any storage but that's OK
                             ! for counting.

    sr%ind => tp%ind
    sr%rearrange = tp%rearrange

    if (sr%rearrange) then
      sr%rdata => tp%rearranged_data
    else
      sr%rdata => tp%the_data
    end if
    sr%dimen = tp%dimen

    !
    !sr%dsl = huge(sr%dsl)    ! set to huge positive values
    !sr%il = -1               ! set to invalid indexes
    !
    sr%overflow = .false.

    call search(tp%root)

    nfound = sr%nfound
  end function kdtree2_r_count_around_point


  subroutine validate_query_storage(n)
    implicit none
    !
    ! make sure we have enough storage for n
    !
    integer, intent(in) :: n

    if ( size(sr%results,1) < n ) then
      call die('KD_TREE_TRANS',&
        'you did not provide enough storage for results(1:n)',1)
    end if
  end subroutine validate_query_storage

  real(kdkind) function square_distance(d,iv,qv) result (res)
    implicit none
    ! distance between iv[1:n] and qv[1:n]
    ! .. Function Return Value ..
    ! re-implemented to improve vectorization.
    ! .. Scalar Arguments ..
    integer :: d
    ! ..
    ! .. Array Arguments ..
    real(kdkind), dimension(:) :: iv, qv
    ! ..
    ! ..
    res = sum( (iv(1:d)-qv(1:d))**2 )
  end function square_distance

  recursive subroutine search(node)
    implicit none
    !
    ! This is the innermost core routine of the kd-tree search.  Along
    ! with "process_terminal_node", it is the performance bottleneck.
    !
    ! This version uses a logically complete secondary search of
    ! "box in bounds", whether the sear
    !
    type(tree_node), pointer :: node
    ! ..
    type(tree_node), pointer :: ncloser, nfarther
    !
    integer :: cut_dim, i
    ! ..
    real(kdkind) :: qval, dis
    real(kdkind) :: ballsize
    real(kdkind), pointer, dimension(:) :: qv
    type(interval), pointer, dimension(:) :: box

    if ( ( associated(node%left) .and. &
           associated(node%right) ) .eqv. .false. ) then
      ! we are on a terminal node
      if ( sr%nn == 0 ) then
        call process_terminal_node_fixedball(node)
      else
        call process_terminal_node(node)
      end if
    else
      ! we are not on a terminal node
      qv => sr%qv(1:)
      cut_dim = node%cut_dim
      qval = qv(cut_dim)

      if ( qval < node%cut_val ) then
        ncloser => node%left
        nfarther => node%right
        dis = (node%cut_val_right - qval)**2
!       extra = node%cut_val - qval
      else
        ncloser => node%right
        nfarther => node%left
        dis = (node%cut_val_left - qval)**2
!       extra = qval- node%cut_val_left
      end if

      if ( associated(ncloser) ) call search(ncloser)

      ! we may need to search the second node.
      if ( associated(nfarther) ) then
        ballsize = sr%ballsize
!       dis = extra**2
        if ( dis <= ballsize ) then
          !
          ! we do this separately as going on the first cut dimen is often
          ! a good idea.
          ! note that if extra**2 < sr%ballsize, then the next
          ! check will also be false.
          !
          box => node%box(1:)
          do i = 1, sr%dimen
            if ( i /= cut_dim ) then
              dis = dis + dis2_from_bnd(qv(i),box(i)%lower,box(i)%upper)
              if ( dis > ballsize ) then
                return
              end if
            end if
          end do

          !
          ! if we are still here then we need to search mroe.
          !
          call search(nfarther)
        end if
      end if
    end if
  end subroutine search

  real(kdkind) function dis2_from_bnd(x,amin,amax) result (res)
    implicit none
    real(kdkind), intent(in) :: x, amin, amax

    if ( x > amax ) then
      res = (x-amax)**2;
      return
    else
      if ( x < amin ) then
        res = (amin-x)**2;
        return
      else
        res = 0.0
        return
      end if
    end if
  end function dis2_from_bnd

!  logical function box_in_search_range(node, sr) result(res)
!    implicit none
!    !
!    ! Return the distance from 'qv' to the CLOSEST corner of node's
!    ! bounding box
!    ! for all coordinates outside the box.   Coordinates inside the box
!    ! contribute nothing to the distance.
!    !
!    type(tree_node), pointer :: node
!    type(tree_search_record), pointer :: sr
!
!    integer :: dimen, i
!    real(kdkind) :: dis, ballsize
!    real(kdkind) :: l, u
!
!    dimen = sr%dimen
!    ballsize = sr%ballsize
!    dis = 0.0
!    res = .true.
!    do i = 1, dimen
!      l = node%box(i)%lower
!      u = node%box(i)%upper
!      dis = dis + (dis2_from_bnd(sr%qv(i),l,u))
!      if ( dis > ballsize ) then
!        res = .false.
!        return
!      end if
!    end do
!    res = .true.
!  end function box_in_search_range

  subroutine process_terminal_node(node)
    implicit none
    !
    ! Look for actual near neighbors in 'node', and update
    ! the search results on the sr data structure.
    !
    type(tree_node), pointer :: node
    !
    real(kdkind), pointer, dimension(:) :: qv
    integer, pointer, dimension(:) :: ind
    real(kdkind), pointer, contiguous, dimension(:,:) :: rdata
    !
    integer :: dimen, i, indexofi, k, centeridx, correltime
    real(kdkind) :: ballsize, sd, newpri
    logical :: rearrange
    type(pq), pointer :: pqp
    !
    ! copy values from sr to local variables
    !
    !
    ! Notice, making local pointers with an EXPLICIT lower bound
    ! seems to generate faster code.
    ! why?  I don't know.
    qv => sr%qv(1:)
    pqp => sr%pq
    dimen = sr%dimen
    ballsize = sr%ballsize
    rearrange = sr%rearrange
    ind => sr%ind(1:)
    rdata => sr%rdata
    centeridx = sr%centeridx
    correltime = sr%correltime

    !    doing_correl = (centeridx >= 0)  ! Do we have a decorrelation window?
    !    include_point = .true.    ! by default include all points
    ! search through terminal bucket.

    mainloop: &
    do i = node%l, node%u
      if ( rearrange ) then
        sd = 0.0
        do k = 1, dimen
          sd = sd + (rdata(k,i) - qv(k))**2
          if ( sd > ballsize ) cycle mainloop
        end do
        indexofi = ind(i)  ! only read it if we have not broken out
      else
        indexofi = ind(i)
        sd = 0.0
        do k = 1, dimen
          sd = sd + (rdata(k,indexofi) - qv(k))**2
          if ( sd > ballsize ) cycle mainloop
        end do
      end if

      if ( centeridx > 0 ) then ! doing correlation interval?
        if ( abs(indexofi-centeridx) < correltime ) cycle mainloop
      end if

      !
      ! two choices for any point.  The list so far is either undersized,
      ! or it is not.
      !
      ! If it is undersized, then add the point and its distance
      ! unconditionally.  If the point added fills up the working
      ! list then set the sr%ballsize, maximum distance bound
      ! (largest distance on list) to be that distance, instead of
      ! the initialized +infinity.
      !
      ! If the running list is full size, then compute the
      ! distance but break out immediately if it is larger
      ! than sr%ballsize, "best squared distance" (of the largest element),
      ! as it cannot be a good neighbor.
      !
      ! Once computed, compare to best_square distance.
      ! if it is smaller, then delete the previous largest
      ! element and add the new one.

      if ( sr%nfound < sr%nn ) then
        !
        ! add this point unconditionally to fill list.
        !
        sr%nfound = sr%nfound+1
        newpri = pq_insert(pqp,sd,indexofi)
        if ( sr%nfound == sr%nn ) ballsize = newpri
        ! we have just filled the working list.
        ! put the best square distance to the maximum value
        ! on the list, which is extractable from the PQ.

      else
        !
        ! now, if we get here,
        ! we know that the current node has a squared
        ! distance smaller than the largest one on the list, and
        ! belongs on the list.
        ! Hence we replace that with the current one.
        !
        ballsize = pq_replace_max(pqp,sd,indexofi)
      end if
    end do mainloop
    !
    ! Reset sr variables which may have changed during loop
    !
    sr%ballsize = ballsize
  end subroutine process_terminal_node

  subroutine process_terminal_node_fixedball(node)
    implicit none
    !
    ! Look for actual near neighbors in 'node', and update
    ! the search results on the sr data structure, i.e.
    ! save all within a fixed ball.
    !
    type (tree_node), pointer :: node
    !
    real(kdkind), pointer, dimension(:) :: qv
    integer, pointer, dimension(:) :: ind
    real(kdkind), pointer, contiguous, dimension(:,:) :: rdata
    !
    integer :: nfound
    integer :: dimen, i, indexofi, k
    integer :: centeridx, correltime, nn
    real(kdkind) :: ballsize, sd
    logical :: rearrange
    !
    ! copy values from sr to local variables
    !
    qv => sr%qv(1:)
    dimen = sr%dimen
    ballsize = sr%ballsize
    rearrange = sr%rearrange
    ind => sr%ind
    rdata => sr%rdata
    centeridx = sr%centeridx
    correltime = sr%correltime
    nn = sr%nn ! number to search for
    nfound = sr%nfound

    ! search through terminal bucket.
    mainloop: &
    do i = node%l, node%u
      !
      ! two choices for any point.  The list so far is either undersized,
      ! or it is not.
      !
      ! If it is undersized, then add the point and its distance
      ! unconditionally.  If the point added fills up the working
      ! list then set the sr%ballsize, maximum distance bound
      ! (largest distance on list) to be that distance, instead
      ! of the initialized +infinity.
      !
      ! If the running list is full size, then compute the
      ! distance but break out immediately if it is larger
      ! than sr%ballsize, "best squared distance" (of the largest element),
      ! as it cannot be a good neighbor.
      !
      ! Once computed, compare to best_square distance.
      ! if it is smaller, then delete the previous largest
      ! element and add the new one.

      ! which index to the point do we use?

      if ( rearrange ) then
        sd = 0.0
        do k = 1, dimen
          sd = sd + (rdata(k,i) - qv(k))**2
          if ( sd > ballsize ) cycle mainloop
        end do
        indexofi = ind(i)  ! only read it if we have not broken out
      else
        indexofi = ind(i)
        sd = 0.0
        do k = 1, dimen
          sd = sd + (rdata(k,indexofi) - qv(k))**2
          if ( sd > ballsize ) cycle mainloop
        end do
      end if

      if ( centeridx > 0 ) then ! doing correlation interval?
        if ( abs(indexofi-centeridx) < correltime ) cycle mainloop
      end if

      nfound = nfound+1
      if ( nfound > sr%nalloc ) then
        ! oh nuts, we have to add another one to the tree but
        ! there isn't enough room.
        sr%overflow = .true.
      else
        sr%results(nfound)%dis = sd
        sr%results(nfound)%idx = indexofi
      end if
    end do mainloop
    !
    ! Reset sr variables which may have changed during loop
    !
    sr%nfound = nfound
  end subroutine process_terminal_node_fixedball

  subroutine kdtree2_n_nearest_brute_force(tp,qv,nn,results)
    implicit none
    ! find the 'n' nearest neighbors to 'qv' by exhaustive search.
    ! only use this subroutine for testing, as it is SLOW!  The
    ! whole point of a k-d tree is to avoid doing what this subroutine
    ! does.
    type(kdtree2), pointer :: tp
    real(kdkind), dimension(:), intent(in) :: qv
    integer, intent(in) :: nn
    type(kdtree2_result), dimension(:) :: results

    integer :: i, j, k
    real(kdkind), allocatable, dimension(:) :: all_distances
    ! ..
    allocate(all_distances(tp%n))
    do i = 1, tp%n
      all_distances(i) = square_distance(tp%dimen,qv,tp%the_data(:,i))
    end do
    ! now find 'n' smallest distances
    do i = 1, nn
      results(i)%dis = huge(1.0)
      results(i)%idx = -1
    end do
    do i = 1, tp%n
      if ( all_distances(i) < results(nn)%dis ) then
        ! insert it somewhere on the list
        do j = 1, nn
          if ( all_distances(i) < results(j)%dis ) exit
        end do
        ! now we know 'j'
        do k = nn - 1, j, -1
          results(k+1) = results(k)
        end do
        results(j)%dis = all_distances(i)
        results(j)%idx = i
      end if
    end do
    deallocate(all_distances)
  end subroutine kdtree2_n_nearest_brute_force

  subroutine kdtree2_r_nearest_brute_force(tp,qv,r2,nfound,results)
    implicit none
    ! find the nearest neighbors to 'qv' with distance**2 <= r2 by
    ! exhaustive search.
    ! only use this subroutine for testing, as it is SLOW!  The
    ! whole point of a k-d tree is to avoid doing what this subroutine
    ! does.
    type(kdtree2), pointer :: tp
    real(kdkind), dimension(:), intent(in) :: qv
    real(kdkind), intent(in) :: r2
    integer, intent(out) :: nfound
    type(kdtree2_result), dimension(:) :: results

    integer :: i, nalloc
    real(kdkind), dimension(:), allocatable :: all_distances
    ! ..
    allocate(all_distances(tp%n))
    do i = 1, tp%n
      all_distances(i) = square_distance(tp%dimen,qv,tp%the_data(:,i))
    end do

    nfound = 0
    nalloc = size(results,1)

    do i = 1, tp%n
      if ( all_distances(i) < r2 ) then
        ! insert it somewhere on the list
        if ( nfound < nalloc ) then
          nfound = nfound+1
          results(nfound)%dis = all_distances(i)
          results(nfound)%idx = i
        end if
      end if
    end do
    deallocate (all_distances)
    call kdtree2_sort_results(nfound,results)
  end subroutine kdtree2_r_nearest_brute_force

  subroutine kdtree2_sort_results(nfound,results)
    implicit none
    !  Use after search to sort results(1:nfound) in order of increasing
    !  distance.
    integer, intent(in) :: nfound
    type(kdtree2_result), dimension(:), target :: results
    !
    if ( nfound > 1 ) call heapsort_struct(results,nfound)
  end subroutine kdtree2_sort_results

!  subroutine heapsort(a,ind,n)
!    implicit none
!    !
!    ! Sort a(1:n) in ascending order, permuting ind(1:n) similarly.
!    !
!    ! If ind(k) = k upon input, then it will give a sort index upon output.
!    !
!    integer, intent(in) :: n
!    real(kdkind), dimension(:), intent(inout) :: a
!    integer, dimension(:), intent(inout) :: ind
!    !
!    real(kdkind) :: rval   ! temporary for a value from a()
!    integer :: ival  ! temporary for a value from ind()
!
!    integer :: i, j
!    integer :: ileft, iright
!
!    ileft = n/2+1
!    iright = n
!
!    !    do i=1,n
!    !       ind(i)=i
!    ! Generate initial idum array
!    !    end do
!
!    if ( n == 1 ) return
!
!    do
!      if ( ileft > 1 ) then
!        ileft = ileft-1
!        rval = a(ileft)
!        ival = ind(ileft)
!      else
!        rval = a(iright)
!        ival = ind(iright)
!        a(iright) = a(1)
!        ind(iright) = ind(1)
!        iright = iright-1
!        if ( iright == 1 ) then
!          a(1) = rval
!          ind(1) = ival
!          return
!        end if
!      end if
!      i = ileft
!      j = 2*ileft
!      do while ( j <= iright )
!        if ( j < iright ) then
!          if ( a(j) < a(j+1) ) j = j+1
!        end if
!        if ( rval < a(j) ) then
!          a(i) = a(j)
!          ind(i) = ind(j)
!          i = j
!          j = j+j
!        else
!          j = iright+1
!        end if
!      end do
!      a(i) = rval
!      ind(i) = ival
!    end do
!  end subroutine heapsort

  subroutine heapsort_struct(a,n)
    implicit none
    !
    ! Sort a(1:n) in ascending order
    !
    integer, intent(in) :: n
    type(kdtree2_result), dimension(:), intent(inout) :: a
    !
    type(kdtree2_result) :: rval ! temporary value

    integer :: i, j
    integer :: ileft, iright

    ileft = n/2+1
    iright = n

    !    do i=1,n
    !       ind(i)=i
    ! Generate initial idum array
    !    end do

    if ( n == 1 ) return

    do
      if ( ileft > 1 ) then
        ileft = ileft-1
        rval = a(ileft)
      else
        rval = a(iright)
        a(iright) = a(1)
        iright = iright-1
        if ( iright == 1 ) then
          a(1) = rval
          return
        end if
      end if
      i = ileft
      j = 2*ileft
      do while ( j <= iright )
        if ( j < iright ) then
          if ( a(j)%dis < a(j+1)%dis ) j = j+1
        end if
        if ( rval%dis < a(j)%dis ) then
          a(i) = a(j)
          i = j
          j = j+j
        else
          j = iright+1
        end if
      end do
      a(i) = rval
    end do
  end subroutine heapsort_struct

end module mod_kdtree2

! vim: tabstop=8 expandtab shiftwidth=2 softtabstop=2
